import numpy as np
import pygame
import time
import random

SCALE = 0.4  # 最小分辨率，0.5mm
NET_WIDTH = 3    #  mm
NET_SIZE = 40  # mm 外围

RADIOS = 15  # mm 铁块半径
CILCLE_INTERVAL = 40  # mm 铁块间距

_NUM_COL = 3 # 铁块布置数量
_NUM_RAW = 3

_WIDTH = 160 # mm
_HEIGHT = 160 # mm
DISPLAY = True
LOG = True
_PATH = "./log.csv"


def gen_log(data:list):
    with open(_PATH, 'a', encoding="utf-8") as f:
        instr = ""
        for i in data:
            instr += f"{i:.3f} "
        f.write(instr[:-1]+"\n")

def surf_to_ndarray(surf:pygame.Surface):
    """pygame image to cv2 form image"""
    frame_np = np.array(pygame.surfarray.pixels3d(surf))
    frame_np = np.transpose(frame_np, (1, 0, 2))
    return frame_np

def ndarray_to_surf(ndarray:np.ndarray):
    conv_img = pygame.surfarray.make_surface(ndarray)
    conv_img = pygame.transform.rotate(conv_img, 90)
    return pygame.transform.flip(conv_img,False,True)

def _gen_points():
    points_col = [[[0, 0], [0, int(_HEIGHT//SCALE)]], ]
    points_row = [[[0, 0], [int(_WIDTH//SCALE), 0]],  ]
    while True:
        newx = points_col[-1][0][0] + int(NET_SIZE//SCALE)
        newy = points_row[-1][0][1] + int(NET_SIZE//SCALE)
        if newx < _WIDTH//SCALE:
            points_col.append([[newx, 0], [newx, int(_HEIGHT//SCALE)]])
        if newy < _HEIGHT//SCALE:
            points_row.append([[0, newy], [int(_WIDTH//SCALE), newy]])
        
        if newx >= _WIDTH//SCALE and newy >= _HEIGHT//SCALE:
            break
        # time.sleep(1)
    return points_col, points_row

def draw_net(surf):
    pc, pr = _gen_points()
    for pair in pc:
        pygame.draw.line(surf, (0, 255, 0), pair[0], pair[1], int(NET_WIDTH//SCALE))
        # pygame.draw.line(surf_world, (0, 0, 0), False, pc, int(NET_WIDTH//SCALE))
    for pair in pr:
        pygame.draw.line(surf, (0, 255, 0), pair[0], pair[1], int(NET_WIDTH//SCALE))


def draw_block(surf, x, y):
    list_block =[]
    for i in range(_NUM_COL):
        for j in range(_NUM_RAW):
            now_x = x+i*(CILCLE_INTERVAL//SCALE)
            now_y = y+j*(CILCLE_INTERVAL//SCALE)
            pygame.draw.circle(surf, (255, 0, 0), (now_x, now_y), (RADIOS//SCALE))
            minx = now_x-(RADIOS//SCALE) if now_x-(RADIOS//SCALE) > 0 else 0
            maxx = now_x+(RADIOS//SCALE) if now_x+(RADIOS//SCALE) < _WIDTH//SCALE else _WIDTH//SCALE
            miny = now_y-(RADIOS//SCALE) if now_y-(RADIOS//SCALE) > 0 else 0
            maxy = now_y+(RADIOS//SCALE) if now_y+(RADIOS//SCALE) < _HEIGHT//SCALE else _HEIGHT//SCALE
            list_block.append([  int(minx), 
                                 int(maxx),
                                 int(miny), 
                                 int(maxy)])
    return list_block


def calc_affect(array:np.ndarray, points:list,):
    """计算每个铁块被遮挡的比例，输出这些比例"""
    temp_array = np.zeros((int(2*RADIOS//SCALE), int(2*RADIOS//SCALE)), dtype=np.uint16)
    affect_rates = []
    round_area = (np.pi*((RADIOS//SCALE)**2))
    for _, i in enumerate(points):
        minx, maxx, miny, maxy = i
        
        for channel in range(3):
            temp_array[0:maxy-miny,0:maxx-minx] = temp_array[0:maxy-miny,0:maxx-minx] + array[miny:maxy,minx:maxx,  channel]
        
        num_affect = len(temp_array[temp_array>255]) 
        affect_rates.append(num_affect/round_area)
        
        temp_array = np.zeros((int(2*RADIOS//SCALE), int(2*RADIOS//SCALE)), dtype=np.uint16)
    return affect_rates


def update(surf_w, surf_c):
    surf_c.fill((0, 0, 0))
    x = random.randint(RADIOS, _WIDTH-RADIOS)
    y = random.randint(RADIOS, _HEIGHT-RADIOS)
    list_block = draw_block(surf_c, x, y)
    img_world = surf_to_ndarray(surf_w)
    img_circle = surf_to_ndarray(surf_c)
    
    out = img_world+img_circle
    # print(calc_affect(out, list_block))
    surf_o = ndarray_to_surf(out)
    return surf_o, calc_affect(out, list_block)


def main():
    clock = pygame.time.Clock()
    pygame.init()
    
    screen = pygame.display.set_mode((_WIDTH//SCALE, _HEIGHT//SCALE)) if DISPLAY else None
    
    surf_world = pygame.Surface((_WIDTH//SCALE, _HEIGHT//SCALE))
    surf_circle = pygame.Surface((_WIDTH//SCALE, _HEIGHT//SCALE))
    draw_net(surf_world)
    surf_out, rates = update(surf_world, surf_circle)

    while True:
        surf_out, rates = update(surf_world, surf_circle)
        print(rates)
        if DISPLAY:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    return
            
            screen.blit(surf_out, (0, 0))
            pygame.display.flip()
        if LOG:
            gen_log(rates)
        # 刷新率
        clock.tick(5)


if __name__ == "__main__":
    main()